#include "../../include/utils/sm2_key_exchange.h"
#include <string.h>
#include <gmssl/sm2.h>
#include <gmssl/sm3.h>
#include <gmssl/rand.h>
#include <gmssl/error.h>

// 安全清零函数
static void secure_zero(void *ptr, size_t size)
{
    volatile unsigned char *p = ptr;
    while (size--)
        *p++ = 0;
}

int sm2_key_exchange_init(SM2_KEY_EXCHANGE_CTX *ctx)
{
    if (!ctx)
        return 0;

    // 初始化上下文成员
    memset(ctx, 0, sizeof(SM2_KEY_EXCHANGE_CTX));
    ctx->has_private_key = 0;

    // 生成密钥对
    if (sm2_key_generate(&ctx->key_pair) != 1)
    {
        return 0;
    }

    ctx->has_private_key = 1;
    return 1;
}

int sm2_generate_exchange_params(SM2_KEY_EXCHANGE_CTX *ctx)
{
    if (!ctx || !ctx->has_private_key)
        return 0;

    // 生成随机数
    if (rand_bytes(ctx->r_a, sizeof(ctx->r_a)) != 1)
    {
        return 0;
    }

    return 1;
}

int sm2_export_public_key(SM2_KEY_EXCHANGE_CTX *ctx, unsigned char *out, size_t *outlen)
{
    if (!ctx || !out || !outlen || !ctx->has_private_key)
        return 0;

    // 直接访问 SM2_KEY 结构体中的 public_key 成员（SM2_POINT 类型）
    // 检查输出缓冲区大小
    if (*outlen < 64) // X坐标和Y坐标各32字节
    {
        *outlen = 64;
        return 0;
    }

    // 复制公钥坐标
    memcpy(out, ctx->key_pair.public_key.x, 32);
    memcpy(out + 32, ctx->key_pair.public_key.y, 32);
    *outlen = 64;

    return 1;
}

// SM3 KDF实现
static int sm3_kdf(const unsigned char *z, size_t z_len,
                   unsigned char *key, size_t key_len)
{
    SM3_CTX ctx;
    unsigned char digest[SM3_DIGEST_SIZE];
    unsigned int i, rounds;
    unsigned int counter = 1;
    size_t offset = 0;

    if (!z || !key || !key_len)
        return 0;

    rounds = (key_len + SM3_DIGEST_SIZE - 1) / SM3_DIGEST_SIZE;

    for (i = 0; i < rounds; i++)
    {
        sm3_init(&ctx);
        sm3_update(&ctx, z, z_len);

        // 添加计数器（大端序）
        unsigned char counter_bytes[4];
        counter_bytes[0] = (counter >> 24) & 0xFF;
        counter_bytes[1] = (counter >> 16) & 0xFF;
        counter_bytes[2] = (counter >> 8) & 0xFF;
        counter_bytes[3] = counter & 0xFF;

        sm3_update(&ctx, counter_bytes, sizeof(counter_bytes));
        sm3_finish(&ctx, digest);

        size_t copy_len = (key_len - offset) > SM3_DIGEST_SIZE ? SM3_DIGEST_SIZE : (key_len - offset);
        memcpy(key + offset, digest, copy_len);

        offset += copy_len;
        counter++;

        if (offset >= key_len)
            break;
    }

    // 安全清除上下文和临时摘要
    secure_zero(&ctx, sizeof(SM3_CTX));
    secure_zero(digest, sizeof(digest));

    return 1;
}

int sm2_key_negotiate(SM2_KEY_EXCHANGE_CTX *initiator_ctx, SM2_KEY_EXCHANGE_CTX *responder_ctx,
                      const unsigned char *remote_pub_x, size_t remote_pub_x_len,
                      const unsigned char *remote_pub_y, size_t remote_pub_y_len,
                      unsigned char *k, size_t klen)
{
    if (!initiator_ctx || !remote_pub_x || !remote_pub_y || !k || klen == 0 ||
        remote_pub_x_len < 32 || remote_pub_y_len < 32 || !initiator_ctx->has_private_key)
    {
        return 0;
    }

    int ret = 0;
    SM2_POINT remote_point;
    SM2_POINT shared_point;
    unsigned char z[64]; // 用于存储共享点的x和y坐标

    // 设置对方公钥点
    memcpy(remote_point.x, remote_pub_x, 32);
    memcpy(remote_point.y, remote_pub_y, 32);

    // 计算共享点
    if (sm2_do_ecdh(&initiator_ctx->key_pair, &remote_point, &shared_point) != 1)
    {
        goto cleanup;
    }

    // 将共享点的坐标拼接作为KDF输入
    memcpy(z, shared_point.x, 32);
    memcpy(z + 32, shared_point.y, 32);

    // 使用KDF生成共享密钥
    if (sm3_kdf(z, sizeof(z), k, klen) != 1)
    {
        goto cleanup;
    }

    ret = 1;

cleanup:
    // 安全清除敏感数据
    secure_zero(&shared_point, sizeof(shared_point));
    secure_zero(z, sizeof(z));

    return ret;
}

void sm2_key_exchange_cleanup(SM2_KEY_EXCHANGE_CTX *ctx)
{
    if (ctx)
    {
        // 清除敏感数据
        if (ctx->has_private_key)
        {
            secure_zero(&ctx->key_pair, sizeof(SM2_KEY));
            ctx->has_private_key = 0;
        }

        secure_zero(ctx->r_a, sizeof(ctx->r_a));
        secure_zero(ctx->r_b, sizeof(ctx->r_b));
        secure_zero(ctx->k, sizeof(ctx->k));
    }
}